use chrono::NaiveDateTime;
use polars::prelude::*;
use pyo3::prelude::*;
use pyo3::types::PyDict;
use pyo3_polars::PyDataFrame;
use std::collections::HashMap;

mod data_load;

#[pyfunction]
fn get_data_demo() -> PyResult<PyDataFrame> {
    let dates = vec![
        NaiveDateTime::parse_from_str("2024-01-01 09:30:00", "%Y-%m-%d %H:%M:%S").unwrap(),
        NaiveDateTime::parse_from_str("2024-01-01 09:31:00", "%Y-%m-%d %H:%M:%S").unwrap(),
        NaiveDateTime::parse_from_str("2024-01-01 09:32:00", "%Y-%m-%d %H:%M:%S").unwrap(),
        NaiveDateTime::parse_from_str("2024-01-01 09:33:00", "%Y-%m-%d %H:%M:%S").unwrap(),
        NaiveDateTime::parse_from_str("2024-01-01 09:34:00", "%Y-%m-%d %H:%M:%S").unwrap(),
        NaiveDateTime::parse_from_str("2024-01-01 09:35:00", "%Y-%m-%d %H:%M:%S").unwrap(),
        NaiveDateTime::parse_from_str("2024-01-01 09:36:00", "%Y-%m-%d %H:%M:%S").unwrap(),
        NaiveDateTime::parse_from_str("2024-01-01 09:37:00", "%Y-%m-%d %H:%M:%S").unwrap(),
        NaiveDateTime::parse_from_str("2024-01-01 09:38:00", "%Y-%m-%d %H:%M:%S").unwrap(),
        NaiveDateTime::parse_from_str("2024-01-01 09:39:00", "%Y-%m-%d %H:%M:%S").unwrap(),
    ];
    // 创建时间戳列(毫秒时间戳)
    let timestamps: Vec<i64> = dates
        .iter()
        .map(|dt| dt.and_utc().timestamp_millis())
        .collect();
    let opens = vec![
        100.0, 101.5, 102.3, 101.8, 102.5, 103.0, 102.7, 103.2, 103.5, 104.0,
    ];
    let highs = vec![
        101.0, 102.8, 103.0, 102.5, 103.2, 103.8, 103.5, 104.0, 104.2, 104.5,
    ];
    let lows = vec![
        99.5, 101.2, 101.8, 101.3, 102.0, 102.5, 102.3, 102.8, 103.0, 103.5,
    ];
    let closes = vec![
        100.8, 102.5, 102.0, 102.2, 103.0, 103.2, 103.0, 103.8, 104.0, 104.2,
    ];
    let volumes = vec![
        1000.0, 1500.0, 1200.0, 1800.0, 2000.0, 1700.0, 1900.0, 2100.0, 1600.0, 2200.0,
    ];
    // 构建Polars DataFrame
    let df = df! [
        "timestamp" => timestamps,
        "datetime" => dates,
        "open" => opens,
        "high" => highs,
        "low" => lows,
        "close" => closes,
        "volume" => volumes,
    ]
        .map_err(|e| {
            PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(format!("Polars error: {}", e))
        })?;
    // 包装DataFrame并返回给Python
    Ok(PyDataFrame(df))
}
async fn test_load_data(
    table_names: Vec<String>,
    s_date: &str,
    e_date: &str,
) -> HashMap<String, DataFrame> {
    let result = data_load::load_data_parallel_optimized(table_names, s_date, e_date).await;
    result.unwrap()
}
#[pyfunction]
fn get_data(
    base_keys: Vec<String>,
    trade_keys: Vec<String>,
    s_date: &str,
    e_date: &str,
) -> PyResult<(PyDataFrame, Py<PyDict>)> {
    // 回测参数
    //let base_keys = vec!["CFFEX_IF00_60"];
    //let trade_keys = vec!["CFFEX_IF00_60", "CFFEX_IC00_60"];
    //let s_date = "20250101";
    //let e_date = "20250201";
    let rt = tokio::runtime::Builder::new_multi_thread()
        .enable_all()
        .build()
        .unwrap();
    let (data_map, data_map2) = rt.block_on(async {
        tokio::join!(
            test_load_data(base_keys.clone(), s_date, e_date),
            test_load_data(trade_keys.clone(), s_date, e_date)
        )
    });
    // 获取第一个 DataFrame
    let base_df = data_map
        .get(&base_keys[0])
        .ok_or_else(|| PyErr::new::<pyo3::exceptions::PyKeyError, _>("Base key not found"))?
        .clone();

    // 转换 HashMap<String, DataFrame> 为 Python 字典
    Python::with_gil(|py| {
        let result_dict = PyDict::new(py);

        for (key, df) in data_map2 {
            result_dict.set_item(key, PyDataFrame(df))?;
        }

        Ok((PyDataFrame(base_df), result_dict.into()))
    })
}

#[pymodule]
fn data(m: &Bound<'_, PyModule>) -> PyResult<()> {
    m.add_function(wrap_pyfunction!(get_data_demo, m)?)?;
    m.add_function(wrap_pyfunction!(get_data, m)?)?;
    Ok(())
}
#[cfg(test)]
mod tests {
    use super::*;
    use tokio;

    #[tokio::test]
    async fn test1() {
        let table_names = vec!["SHFE_au00_60".to_owned()];
        let result = test_load_data(table_names, "20250101", "20250103").await;

        println!("成功: {:?}", result);
        // 可以添加一些断言
        assert!(!result.is_empty());
    }
}
